package com.biwow.ai.helloworld.service;

import io.grpc.Status;
import io.grpc.StatusRuntimeException;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.QdrantGrpcClient;
import io.qdrant.client.grpc.Collections;
import io.qdrant.client.grpc.Points.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutionException;

import static io.qdrant.client.PointIdFactory.id;
import static io.qdrant.client.QueryFactory.nearest;
import static io.qdrant.client.ValueFactory.value;
import static io.qdrant.client.VectorsFactory.vectors;
import static io.qdrant.client.WithPayloadSelectorFactory.enable;

@Service
public class QdrantService {
    private final QdrantClient qdrantClient;

    @Autowired
    public QdrantService(
            @Value("${spring.qdrant.host}") String host,
            @Value("${spring.qdrant.port}") int port) {
        try {
            this.qdrantClient = new QdrantClient(
                    QdrantGrpcClient.newBuilder(host, port, false).build()
            );
        } catch (Exception e) {
            throw new IllegalStateException("Qdrant客户端初始化失败", e);
        }
    }

    /**
     * 创建集合
     *
     * @param collectionName 指定要创建的集合名称，需保证唯一性
     * @param dim            定义向量维度，同一集合中所有向量必须维度相同
     */
    public void createCollection(String collectionName, int dim) {
        qdrantClient.createCollectionAsync(collectionName,
                Collections.VectorParams.newBuilder()
                        .setDistance(Collections.Distance.Cosine)
                        .setSize(dim)
                        .build());
    }

    public boolean existsCollection(String collectionName) {
        try {
            qdrantClient.getCollectionInfoAsync(collectionName).get();
            return true; // 成功获取信息说明集合存在
        } catch (Exception e) {
            return false;
        }
    }

    /**
     * 插入向量记录
     *
     * @param collectionName 目标集合名称，需确保已存在
     * @param pointId        点的唯一标识符
     * @param vector         浮点数数组表示的向量数据
     * @param payload        键值对形式的元数据，支持动态字段
     */
    public UpdateResult insertRecord(String collectionName, long pointId, float[] vector, Map<String, Object> payload) throws Exception {
        PointStruct point = PointStruct.newBuilder()
                .setId(id(pointId))
                .setVectors(vectors(vector))
                .putAllPayload(payload != null ?
                        payload.entrySet().stream()
                                .collect(java.util.stream.Collectors.toMap(
                                        Map.Entry::getKey,
                                        e -> value((String) e.getValue())))
                        : Map.of())
                .build();

        // 执行插入操作
        return qdrantClient.upsertAsync(collectionName, List.of(point)).get();
    }

    /**
     * @param collectionName 目标集合名称，需确保已存在
     * @param limit          限制返回结果数量
     * @param vector         查询向量，维度需与集合定义一致
     * @param filters        过滤参数
     */
    public List<ScoredPoint> searchSimilarRecords(
            String collectionName, int limit, float[] vector, Map<String, Object> filters) throws Exception {

        // 构建查询条件
        QueryPoints.Builder queryBuilder = QueryPoints.newBuilder()
                .setCollectionName(collectionName)
                .setLimit(limit)
                .setQuery(nearest(vector))
                .setWithPayload(enable(true));

        // 添加过滤条件
        if (filters != null && !filters.isEmpty()) {
            Filter.Builder filterBuilder = Filter.newBuilder();
            filters.forEach((field, value) -> {
                FieldCondition.Builder fieldCondition = FieldCondition.newBuilder()
                        .setKey(field);

                if (value instanceof Map<?, ?> range) {
                    // 处理范围查询
                    Range.Builder rangeBuilder = Range.newBuilder();
                    if (range.containsKey("gte"))
                        rangeBuilder.setGte(((Number) range.get("gte")).floatValue());
                    if (range.containsKey("lte"))
                        rangeBuilder.setLte(((Number) range.get("lte")).floatValue());
                    fieldCondition.setRange(rangeBuilder);
                } else {
                    // 处理精确匹配
                    fieldCondition.setMatch(Match.newBuilder()
                            .setKeyword(value.toString()));
                }
                filterBuilder.addMust(Condition.newBuilder()
                        .setField(fieldCondition));
            });
            queryBuilder.setFilter(filterBuilder);
        }

        // 执行查询
        return qdrantClient.queryAsync(queryBuilder.build()).get();
    }
}
